/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package grex.Data;

import WekaModels.WekaPredictiveModel;
import grex.Prediction;
import grex.PredictionContainer;
import java.util.Arrays;
import java.util.Hashtable;

/**
 *
 * @author RIK
 */
public class ExtractionArffTableModel extends ArffTableModel {

    WekaPredictiveModel model;
    private Hashtable<String, double[][]> extractionTrainFolds = new Hashtable<String, double[][]>();
    private Hashtable<String, double[][]> extractionValFolds = new Hashtable<String, double[][]>();
    private Hashtable<String, double[][]> extractionTestFolds = new Hashtable<String, double[][]>();
    private boolean useOracle = false;

    public ExtractionArffTableModel(ArffTableModel data, WekaPredictiveModel model) {
        super(data);
        this.model = model;
    }

    /*
     * public double[][] getTestFold(String fold) {
     *
     *
     * }
     *
     * public double[][] getValFold(String fold) {
     *
     * return data; }
     */
    //Should always be used to get a fold as lazy initzilation is used.
    public void setUseOracle(boolean useOracle) {
        this.useOracle = useOracle;
    }

    @Override
    public double[][] getTrainFold(String fold) {
        double[][] extData = extractionTrainFolds.get(fold);
        if (extData == null) {
            double[][] data = super.getTrainFold(fold);
            extData = new double[data.length][data[0].length];
            model.initFold(fold);
            model.train();
            PredictionContainer pc = model.getPcTrain();
            if (!useOracle) {
                for (int i = 0; i < data.length; i++) {
                    extData[i] = Arrays.copyOfRange(data[i], 0, data[i].length);
                    Prediction p = pc.get(data[i]);
                    extData[i][extData[i].length - 1] = p.getPrediction();
                }
            } else {
                double[][] testData = super.getTestFold(fold);
                double[][] oracleData = new double[testData.length][testData[0].length];
                model.execute(model.getPcTest());
                pc = model.getPcTest();
                for (int i = 0; i < testData.length; i++) {
                    oracleData[i] = Arrays.copyOfRange(testData[i], 0, testData[i].length);
                    Prediction p = pc.get(testData[i]);
                    oracleData[i][oracleData[i].length - 1] = p.getPrediction();
                }
                extData = oracleData;

            }

            extractionTrainFolds.put(fold, extData);
        }
        return extData;
    }
}
